import wordcloud
import jieba
import imageio.v2 as imageio
from PIL import Image

# 1、读取文件
txt = open("../三国演义.txt", "r", encoding="utf-8").read()

# 2、分词
words = jieba.cut(txt)

# 3、统计
excludes = {"却说", "将军", "主公", "左右", "丞相", "军马", "如此", "不可", "荆州", "军士", "不能", "二人", "商议",
            "于是", "一人", "东吴", "不散", "今日", "次日", "天下", "都督", "大喜", "引兵", "陛下", "魏兵", "汉中",
            "上马", "只见", "蜀兵", "人马", "后主", "众将", "不敢", "不知", "太守", "大叫", "夫人", "此人", "后人",
            "先主", "背后", "城中", "如何"}
counts = {}
for w in words:
    if (len(w) < 2) or (len(w) > 4) or (w in excludes):
        continue
    if w == "皇叔" or w == "玄德" or w == "玄德曰":
        w = "刘备"
    if w == "关公" or w == "云长" or w == "云长曰":
        w = "关羽"
    if w == "翼德":
        w = "张飞"
    if w == "孔明" or w == "卧龙" or w == "孔明曰":
        w = "诸葛亮"
    if w == "孟德" or w == "孟德曰":
        w = "曹操"
    counts[w] = counts.get(w, 0) + 1

# 4、顺序
items = list(counts.items())
items.sort(key=lambda x: x[1], reverse=True)

# 5、结果
top_word = ""
for i in range(15):
    c, count = items[i]
    top_word = top_word + c + " "

# 6、词云显示
mask = imageio.imread("../chapter05/china.jpeg")
# w = wordcloud.WordCloud(width=1000, height=700, font_path="msyh.ttc", mask=mask, min_font_size=10, max_font_size=180)
w = wordcloud.WordCloud(width=1000, height=700, font_path="/System/Library/fonts/PingFang.ttc", mask=mask,
                        background_color="white", min_font_size=10, max_font_size=180)
w.generate(top_word)
w.to_file("sanguo-cloud.png")
a = Image.open("sanguo-cloud.png")
a.show()
